package main

import (
	"ai-gateway/common"
	"ai-gateway/config"
	"ai-gateway/proto"
	"context"
	"fmt"
	"io"
	"net"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
)

type BusinessServer struct {
	proto.UnimplementedBusinessServiceServer
}

// 统一入口：自动判断流式/非流式，业务端只需调用这一个方法
func (s *BusinessServer) UnifiedProcess(req *proto.AIRequest, stream proto.BusinessService_UnifiedProcessServer) error {
	// 1. 判断是否为流式场景（复用配置或数据库逻辑）
	isStream := false
	if config.Config.StreamScenarios != nil {
		isStream = config.Config.StreamScenarios[req.Scenario]
	}

	if isStream {
		common.Logger.Info().Str("scenario", req.Scenario).Msg("自动路由到流式处理")
		// 2. 流式场景：调用原 GenerateContentStream，用适配器兼容接口
		return s.GenerateContentStream(req, &StreamAdapter{stream: stream})
	}

	common.Logger.Info().Str("scenario", req.Scenario).Msg("自动路由到非流式处理")
	// 3. 非流式场景：调用原 GenerateContent，结果通过流返回
	resp, err := s.GenerateContent(stream.Context(), req)
	if err != nil {
		return err
	}
	return stream.Send(resp)
}

type StreamAdapter struct {
	stream proto.BusinessService_UnifiedProcessServer
}

// 实现 context.Context 方法
func (a *StreamAdapter) Context() context.Context {
	return a.stream.Context()
}

// 实现 Send 方法，用于发送业务响应
func (a *StreamAdapter) Send(resp *proto.AIResponse) error {
	return a.stream.Send(resp)
}

// 实现 SendHeader 方法，gRPC 框架可能会调用（透传即可）
func (a *StreamAdapter) SendHeader(md metadata.MD) error {
	return a.stream.SendHeader(md)
}

// 实现 SetHeader 方法，gRPC 框架可能会调用（透传即可）
func (a *StreamAdapter) SetHeader(md metadata.MD) error {
	return a.stream.SetHeader(md)
}

// 实现 SetTrailer 方法，gRPC 框架可能会调用（透传即可）
func (a *StreamAdapter) SetTrailer(md metadata.MD) {
	a.stream.SetTrailer(md)
}

// 实现 SendMsg 方法，gRPC 内部序列化用（透传）
func (a *StreamAdapter) SendMsg(m any) error {
	return a.stream.SendMsg(m)
}

// 实现 RecvMsg 方法，流式场景服务端不接收客户端消息，返回错误即可
func (a *StreamAdapter) RecvMsg(m any) error {
	return fmt.Errorf("UnifiedProcess 服务端流不支持接收客户端消息")
}

// BusinessServer 需实现新增的流式方法
func (s *BusinessServer) GenerateContentStream(req *proto.AIRequest, stream proto.BusinessService_GenerateContentStreamServer) error {
	// 权和流量控制逻辑
	// 鉴权
	authAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.AuthServicePort)
	authConn, err := grpc.Dial(authAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return fmt.Errorf("连接鉴权服务失败: %v", err)
	}
	defer authConn.Close()
	authClient := proto.NewAuthServiceClient(authConn)
	authResp, err := authClient.CheckPermission(stream.Context(), req)
	if err != nil || !authResp.Allowed {
		return fmt.Errorf("鉴权失败: %v", err)
	}

	// 流量控制
	trafficAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.TrafficServicePort)
	trafficConn, err := grpc.Dial(trafficAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return fmt.Errorf("连接流量服务失败: %v", err)
	}
	defer trafficConn.Close()
	trafficClient := proto.NewTrafficServiceClient(trafficConn)
	quotaResp, err := trafficClient.CheckQuota(stream.Context(), req)
	if err != nil || !quotaResp.Allowed {
		return fmt.Errorf("流量控制失败: %s", quotaResp.Message)
	}

	// 调用适配器的流式服务
	adapterAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.AdapterServicePort)
	adapterConn, err := grpc.Dial(adapterAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return fmt.Errorf("连接适配器服务失败: %v", err)
	}
	defer adapterConn.Close()
	adapterClient := proto.NewAdapterServiceClient(adapterConn)

	// 获取适配器的流式响应
	adapterStream, err := adapterClient.StreamCallAI(stream.Context(), req)
	if err != nil {
		return fmt.Errorf("适配器流式调用失败: %v", err)
	}

	// 实时转发流式片段给客户端
	var lastResp *proto.AIResponse // 记录最后一个响应用于监控
	for {
		resp, err := adapterStream.Recv()
		if err == io.EOF {
			break // 流式结束
		}
		if err != nil {
			return fmt.Errorf("接收适配器流失败: %v", err)
		}
		// 转发给客户端
		if err := stream.Send(resp); err != nil {
			return fmt.Errorf("转发给客户端失败: %v", err)
		}
		lastResp = resp
	}

	// 上报监控日志（复用最后一个响应的元数据）
	if lastResp != nil {
		monitorAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.MonitorServicePort)
		monitorConn, err := grpc.Dial(monitorAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
		if err == nil {
			defer monitorConn.Close()
			monitorClient := proto.NewMonitorServiceClient(monitorConn)
			_, _ = monitorClient.ReportLog(stream.Context(), &proto.CallLog{
				BusinessSystem: req.BusinessSystem,
				Provider:       lastResp.Provider,
				TokenUsed:      lastResp.TokenUsed,
				Cost:           lastResp.Cost,
			})
		}
	}

	return nil
}

func (s *BusinessServer) GenerateContent(ctx context.Context, req *proto.AIRequest) (*proto.AIResponse, error) {
	// 调用AuthService鉴权，检查业务系统是否有权限，通过鉴权表进行判断
	authAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.AuthServicePort)
	authConn, err := grpc.Dial(authAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return &proto.AIResponse{}, fmt.Errorf("连接鉴权服务失败: %v", err)
	}
	defer authConn.Close()
	authClient := proto.NewAuthServiceClient(authConn)

	authResp, err := authClient.CheckPermission(ctx, req)
	if err != nil {
		return &proto.AIResponse{}, fmt.Errorf("鉴权服务调用失败: %v", err)
	}
	if !authResp.Allowed {
		return &proto.AIResponse{}, fmt.Errorf("业务系统无调用权限")
	}

	// 调用TrafficService检查流量配额-是否超过每日调用上限
	trafficAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.TrafficServicePort)
	trafficConn, err := grpc.Dial(trafficAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return &proto.AIResponse{}, fmt.Errorf("连接流量服务失败: %v", err)
	}
	defer trafficConn.Close()
	trafficClient := proto.NewTrafficServiceClient(trafficConn)

	quotaResp, err := trafficClient.CheckQuota(ctx, req)
	if err != nil || !quotaResp.Allowed {
		return &proto.AIResponse{}, fmt.Errorf("流量控制失败: %s", quotaResp.Message)
	}

	// 连接适配器服务，根据场景选择流式和同步
	adapterAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.AdapterServicePort)
	adapterConn, err := grpc.Dial(adapterAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return &proto.AIResponse{}, fmt.Errorf("连接适配器服务失败: %v", err)
	}
	defer adapterConn.Close()
	adapterClient := proto.NewAdapterServiceClient(adapterConn)

	// 根据配置中的流式场景列表表判断调用方式
	var aiResp *proto.AIResponse
	// 同步调用逻辑
	syncResp, err := adapterClient.CallAI(ctx, req)
	if err != nil {
		return &proto.AIResponse{}, fmt.Errorf("同步调用AI服务失败: %v", err)
	}
	aiResp = syncResp

	// 调用MonitorService记录日志，记录某一次调用的Token消耗、成本
	monitorAddr := fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.MonitorServicePort)
	monitorConn, err := grpc.Dial(monitorAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		common.Logger.Warn().Err(err).Msg("连接监控服务失败，跳过日志上报")
	} else {
		defer monitorConn.Close()
		monitorClient := proto.NewMonitorServiceClient(monitorConn)
		_, _ = monitorClient.ReportLog(ctx, &proto.CallLog{
			BusinessSystem: req.BusinessSystem,
			Provider:       aiResp.Provider,
			TokenUsed:      aiResp.TokenUsed,
			Cost:           aiResp.Cost,
			Error:          "",
		})
	}

	return aiResp, nil
}

func main() {
	//初始化一个日志
	common.InitLogger("business-service")

	// 启动业务服务（供业务系统调用）
	lis, err := net.Listen("tcp", fmt.Sprintf("%s%d", ":", config.Config.BusinessServicePort))
	if err != nil {
		common.Logger.Fatal().Err(err).Msg("监听端口失败")
	}
	s := grpc.NewServer()
	proto.RegisterBusinessServiceServer(s, &BusinessServer{})
	common.Logger.Info().Msg("Business Service 启动: :8080（业务侧调用入口）")

	if err := s.Serve(lis); err != nil {
		common.Logger.Fatal().Err(err).Msg("服务启动失败")
	}
}
